from typing import TypedDict
from langchain_core.messages import HumanMessage
from langgraph.graph import StateGraph
import logging
from .util import write_file

class NovelNodes:
    def __init__(self, llm):
        self.llm = llm
        self.logger = logging.getLogger(__name__)

    def _call_llm(self, prompt):
        self.logger.info(f"调用LLM，提示词：{prompt}")
        return self.llm.invoke([HumanMessage(content=prompt)]).content

    def generate_initial_outline(self, state):
        prompt = "请生成一个包含5个章节的小说大纲，包含主要人物和剧情走向"
        outline = self._call_llm(prompt)
        return {"outline": outline, "iteration": 0}

    def refine_outline(self, state):
        prompt = f"当前大纲：{state['outline']}\n请优化改进这个小说大纲"
        return {"outline": self._call_llm(prompt), "iteration": state['iteration']+1}

    def write_chapter(self, state):
        prompt = f"根据大纲：{state['outline']}\n撰写章节：{state['current_chapter']}，限500字"
        content = self._call_llm(prompt)
        write_file(f"chapters/{state['current_chapter']}.txt", content)
        return {
            "content": content[:50],
            "chapter": state['current_chapter'],
            "word_count": len(content)
        }

    def final_review(self, state):
        prompt = f"最终审核完整小说内容：{state['outline']}\n请检查逻辑连贯性和文风统一性"
        return {"review": self._call_llm(prompt)}